#!/usr/bin/python3.9
# -*- coding: utf-8 -*-
# @Time    :  2025/8/1 14:15
# @Author  : chenxw
# @Email   : gisfanmachel@gmail.com
# @File    : vgisPredictor.py
# @Descr   : AI预测器
# @Software: PyCharm
import os.path
import uuid

from vgis_infer import utils
from vgis_infer.utils import get_file_extensions, allowed_archive, snowflakeId, add_alg_task, update_alg_progress


# 直接对接遥感AI算法发布接口，不是服务网关接口
# 图片预测器
class VGISImagePredictor:
    # 初始化
    def __init__(self, ai_service_url):
        self.ai_service_url = ai_service_url

    # 对图片进行目标识别(支持单个或多个图片）
    # image_datas里是数组
    def predictor_image(self, image_datas, params):
        # 请求AI服务
        payjson = {}
        payjson["input_path"] = image_datas
        payjson["object_name"] = params["object_name"]
        payjson["alg_model"] = params["alg_model"]
        payjson["conf_value"] = params["conf_value"]
        payjson["iou_value"] = params["iou_value"]
        payjson["source_type"] = "image"
        if "pic_type" not in params:
            params["pic_type"] = "common_pics"
        payjson["pic_type"] = params["pic_type"]
        payjson["alg_result"] = "{}_result.json".format(uuid.uuid4())
        if "is_save_db" in params:
            payjson["is_save_db"] = params["is_save_db"]
            payjson["operate_db_info"]=params["operate_db_info"]

        image_widths = []
        image_heights = []
        # 普通图片
        if params["pic_type"] == "common_pics":
            if params["image_type"] == "url":
                for image_data in image_datas:
                    image_width, image_height = utils.get_size_of_url_image_type(image_data)
                    image_widths.append(image_width)
                    image_heights.append(image_height)
            if params["image_type"] == "base_64":
                payjson["input_path"] = []
                for image_data in image_datas:
                    image_path = os.path.join(os.getcwd(), "tmp", uuid.uuid4().hex + ".jpg")
                    image_path = utils.save_base64_image(image_data, image_path)
                    image_width, image_height = utils.get_size_of_image(image_path)
                    image_widths.append(image_width)
                    image_heights.append(image_height)
                    upload_file_path = utils.upload_file_service(self.ai_service_url + "/upload", image_path)
                    payjson["input_path"].append(upload_file_path)
                    os.remove(image_path)
            if params["image_type"] == "path":
                payjson["input_path"] = []
                for image_data in image_datas:
                    image_width, image_height = utils.get_size_of_image(image_data)
                    image_widths.append(image_width)
                    image_heights.append(image_height)
                    upload_file_path = utils.upload_file_service(self.ai_service_url + "/upload", image_data)
                    payjson["input_path"].append(upload_file_path)
        # 空间图片
        elif params["pic_type"] == "spatial_pic":
            payjson["bands"]=params["bands"]
            payjson["input_path"] = []
            payjson["epsg_value"] = params["epsg_value"]
            extension=get_file_extensions(image_datas, with_dot=False)
            for image_data in image_datas:
                if allowed_archive(extension):
                    upload_analysis_file_path = utils.upload_extract_filezip_service(self.ai_service_url + "/upload", image_data,
                                                                            params["analysis_file_extension"])
                else:
                    upload_analysis_file_path = utils.upload_file_service(self.ai_service_url + "/upload", image_data)
                payjson["input_path"].append(upload_analysis_file_path)

        # 支持对
        # token_value = utils.get_ai_service_token(self.ai_service_url)
        return_results = utils.toggle_ai_service(self.ai_service_url + "/vgis_object_detect", payjson)
        all = {}

        index = 0
        if return_results["success"]:
            all["success"] = True
            all["info"] = "执行成功"
            res = {}
            res["version"] = "1.0"
            res["imageListInfers"] = []
            if params["pic_type"] == "common_pics":
                for result_box_list in return_results["yolo_results"]:
                    res2 = {}
                    res2["imageIndex"] = index
                    res2["imageInferResults"] = []
                    res2["imagePath"] = image_datas[index] if params["image_type"] != "base_64" else None
                    res2["imageData"] = None
                    res2["imageWidth"] = image_widths[index]
                    res2["imageHeight"] = image_heights[index]
                    for result_box in result_box_list:
                        xmin = result_box["xmin"]
                        ymin = result_box["ymin"]
                        xmax = result_box["xmax"]
                        ymax = result_box["ymax"]
                        name = result_box["name"]
                        confidence = result_box["confidence"]
                        res2["imageInferResults"].append(
                            {"id": utils.snowflakeId(), "label": name, "text": utils.get_predict_cname_by_ename(name),
                             "shapes": [
                                 {"shape_type": "hbb",
                                  "points": [[xmin, ymin], [xmin, ymax], [xmax, ymax], [xmax, ymin]]}],
                             "conf": confidence})
                    res["imageListInfers"].append(res2)
                    index += 1
            elif params["pic_type"] == "spatial_pic":
                pic_counts=return_results["pic_counts"]
                for pic_index in range(pic_counts):
                    obj = {}
                    obj["count"] = return_results["shp_counts"][pic_index]
                    # 返回在AI服务器上返回的分析结果shp路径
                    obj["shp_path"] = return_results["shp_paths"][pic_index]
                    # 将shp做后处理并转换geojson
                    payload = {}
                    payload["input_shp_path"]=obj["shp_path"]
                    payload["epsg_value"] = params["epsg_value"]
                    geojson_res=utils.toggle_ai_service(self.ai_service_url + "/handle_shp_geojson", payload)
                    obj["new_shp_path"] = geojson_res["new_shp_path"]
                    obj["all_geojson_path"]=geojson_res["all_geojson_path"]
                    obj["all_geojson_url"] = geojson_res["all_geojson_url"]
                    obj["all_zip_path"] = geojson_res["all_zip_path"]
                    obj["all_zip_url"] = geojson_res["all_zip_url"]
                    payload = {}
                    payload["input_shp_path"] = obj["new_shp_path"]
                    # 发布shp服务
                    publish_res = utils.toggle_ai_service(self.ai_service_url + "/publish_shp_service", payload)
                    obj["topojson_vector_tiles_url"]=publish_res["topojson_vector_tiles_url"]
                    obj["geojson_vector_tiles_url"] = publish_res["geojson_vector_tiles_url"]
                    obj["utfgrid_vector_tiles_url"] = publish_res["utfgrid_vector_tiles_url"]
                    obj["mvt_vector_tiles_url"] = publish_res["mvt_vector_tiles_url"]
                    res["imageListInfers"].append(obj)
            all["data"] = res
        else:
            all["success"] = False
            all["info"] = return_results["info"]
        return all

    # 对图片对进行变化检测(支持多个图片对）
    # image_datas里是对象数组[{"img_before_path":img_before_path,"img_after_path":img_after_path}]
    def change_pair_image(self, image_datas, params):
        # 请求AI服务
        payjson = {}
        payjson["input_path"] = image_datas
        payjson["object_type"] = params["object_type"]
        payjson["alg_model"] = params["alg_model"]
        payjson["alg_result"] = "{}_result.json".format(uuid.uuid4())
        if "is_segement" in params:
            payjson["is_segment"] = params["is_segment"]
            payjson["segment_method"]=params["segment_method"]
        if "sat_seg_method" not in params:
            payjson["sat_seg_method"]="default"
        else:
            payjson["sat_seg_method"]=params["sat_seg_method"]
        if "sat_obj_method" not in params:
            payjson["sat_obj_method"]="objdetection"
        else:
            payjson["sat_obj_method"]=params["sat_obj_method"]
        if "is_save_db" in params:
            payjson["is_save_db"] = params["is_save_db"]
            payjson["operate_db_info"]=params["operate_db_info"]

        payjson["epsg_value"] = params["epsg_value"]
        payjson["alg_result"] = "{}_result.json".format(uuid.uuid4())
        # 针对多个图像对进行变化检测，需要调用多次AI服务
        return_results=[]
        for image_data in image_datas:
            img_before_path=image_data["img_before_path"]
            img_after_path=image_data["img_after_path"]
            if allowed_archive(img_before_path):
                upload_analysis_before_file_path = utils.upload_extract_filezip_service(self.ai_service_url + "/upload", img_before_path,
                                                                        params["analysis_file_extension"])
                upload_analysis_after_file_path = utils.upload_extract_filezip_service(self.ai_service_url + "/upload",
                                                                                        img_after_path,
                                                                                        params[
                                                                                            "analysis_file_extension"])
            else:
                upload_analysis_before_file_path = utils.upload_file_service(self.ai_service_url + "/upload", img_before_path)
                upload_analysis_after_file_path=utils.upload_file_service(self.ai_service_url + "/upload", img_after_path)
            payjson["img_before_path"].append(upload_analysis_before_file_path)
            payjson["img_after_path"].append(upload_analysis_after_file_path)
            # token_value = utils.get_ai_service_token(self.ai_service_url)
            return_result = utils.toggle_ai_service(self.ai_service_url + "/vgis_change_detect", payjson)
            return_results.append(return_result)

        all_result=[]
        each_result = {}

        index = 0
        for result in return_results:
            if result:
                each_result["success"] = True
                each_result["info"] = "执行成功"
                res = {}
                res["version"] = "1.0"
                res["pairimageChangeInfers"] = []

                obj = {}
                obj["count"] = result["count"]
                # 返回在AI服务器上返回的分析结果shp路径
                obj["shp_path"] = result["shp_path"]
                # 将shp做后处理并转换geojson
                payload = {}
                payload["input_shp_path"]=obj["shp_path"]
                payload["epsg_value"] = params["epsg_value"]
                geojson_res=utils.toggle_ai_service(self.ai_service_url + "/handle_shp_geojson", payjson)
                obj["new_shp_path"] = geojson_res["new_shp_path"]
                obj["all_geojson_path"]=geojson_res["all_geojson_path"]
                obj["all_geojson_url"] = geojson_res["all_geojson_url"]
                obj["all_zip_path"] = geojson_res["all_zip_path"]
                obj["all_zip_url"] = geojson_res["all_zip_url"]
                payload = {}
                payload["input_shp_path"] = obj["new_shp_path"]
                # 发布shp服务
                publish_res = utils.toggle_ai_service(self.ai_service_url + "/publish_shp_service", payjson)
                obj["topojson_vector_tiles_url"]=publish_res["topojson_vector_tiles_url"]
                obj["geojson_vector_tiles_url"] = publish_res["geojson_vector_tiles_url"]
                obj["utfgrid_vector_tiles_url"] = publish_res["utfgrid_vector_tiles_url"]
                obj["mvt_vector_tiles_url"] = publish_res["mvt_vector_tiles_url"]
                res["pairimageChangeInfer"]=obj
                each_result["data"] = res
            else:
                each_result["success"] = False
                each_result["info"] = return_results["info"]
            all_result.append(each_result)
        return all_result

# 视频预测器
class VGISVideoPredictor:
    # 初始化
    def __init__(self, ai_service_url):
        self.ai_service_url = ai_service_url

    # 对单个视频进行预测
    def predictor_video(self, video_data, params):
        # 请求AI服务
        payjson = {}
        payjson["input_path"] = video_data
        payjson["object_name"] = params["object_name"]
        payjson["alg_model"] = params["alg_model"]
        payjson["conf_value"] = params["conf_value"]
        payjson["iou_value"] = params["iou_value"]
        payjson["source_type"] = "video"
        payjson["alg_result"] = "{}_result.json".format(uuid.uuid4())

        if params["video_type"] == "base_64":
            payjson["video_type"] = "video_frame_pics"
            payjson["input_path"] = []
            frame_index = 0
            for image_data in video_data:
                image_path = os.path.join(os.getcwd(), uuid.uuid4().hex + ".jpg")
                image_path = utils.save_base64_image(image_data, image_path)
                if frame_index == 0:
                    video_width, video_height = utils.get_size_of_image(image_path)
                upload_file_path = utils.upload_file_service(self.ai_service_url + "/upload", image_path)
                payjson["input_path"].append(upload_file_path)
                os.remove(image_path)

        if params["video_type"] == "path" or params["video_type"] == "url" or params["video_type"] == "stream":
            if params["video_type"] == "stream":
                payjson["video_type"] = "video_stream"
            else:
                payjson["video_type"] = "video_file"
            payjson["input_path"] = []
            for video_file in video_data:
                video_info = utils.get_video_info(video_file)
                video_width = video_info["width"]
                video_height = video_info["height"]
                if params["video_type"] == "path":
                    upload_file_path = utils.upload_file_service(self.ai_service_url + "/upload", video_file)
                    payjson["input_path"].append(upload_file_path)
                else:
                    payjson["input_path"].append(video_file)
        # token_value = utils.get_ai_service_token(self.ai_service_url)
        return_results = utils.toggle_ai_service(self.ai_service_url + "/vgis_object_detect", payjson)

        all = {}

        index = 0
        if return_results["success"]:
            all["success"] = True
            all["info"] = "执行成功"
            res = {}
            res["version"] = "1.0"
            res["videoInfers"] = []
            for result_box_list in return_results["yolo_results"]:
                res2 = {}
                res2["frameIndex"] = index
                res2["imageInferResults"] = []

                for result_box in result_box_list:
                    xmin = result_box["xmin"]
                    ymin = result_box["ymin"]
                    xmax = result_box["xmax"]
                    ymax = result_box["ymax"]
                    name = result_box["name"]
                    confidence = result_box["confidence"]
                    res2["imageInferResults"].append(
                        {"id": utils.snowflakeId(), "lable": name, "text": utils.get_predict_cname_by_ename(name),
                         "shapes": [
                             {"shape_type": "hbb", "points": [[xmin, ymin], [xmin, ymax], [xmax, ymax], [xmax, ymin]]}],
                         "conf": confidence})
                res["videoInfers"].append(res2)
                index += 1
            res["videoPath"] = video_data[0] if params["video_type"] != "base_64" else None
            res["videoData"] = None
            res["videoHeight"] = video_height
            res["videoWidth"] = video_width
            all["data"] = res
        else:
            all["success"] = False
            all["info"] = return_results["info"]
        return all


if __name__ == '__main__':
    # 调试url
    ai_service_url = "http://192.168.31.39:28619"
    # 部署url
    # ai_service_url = "http://192.168.31.39:28616"

    # 测试图像检测器
    image_detector = VGISImagePredictor(ai_service_url)
    # 测试路径空间图片（单张tif）
    image_data = [r"E:\AI\test_data\目标识别\airplane\kj1_TW_2019_KJ_0.5.tif"]
    bands=[[1,2,3]]
    params = {"object_name": "airplane", "alg_model": "sat_hr-vgisyolo_v11n-objdet-yolo_hbb", "conf_value": 0.45,
              "iou_value": 0.7,
              "image_data": image_data, "image_type": "path", "pic_type": "spatial_pic", "epsg_value": 4326,"bands":bands}

    # # 测试路径空间图片（空间图片的压缩包）
    # image_data = [r"E:\AI\test_data\目标识别\airplane\air.zip"]
    # # 压缩包里待分析的空间图片的后缀
    # analysis_file_extension = ".jpg"
    # params = {"object_name": "airplane", "alg_model": "sat_hr-yolo_v11n-objdet-yolo_hbb", "conf_value": 0.45,
    #           "iou_value": 0.7,
    #           "image_data": image_data, "image_type": "path", "pic_type": "spatial_pic",
    #           "analysis_file_extension": analysis_file_extension, "epsg_value": 4326}
    #
    # # 测试路径普通图片(单张)
    # image_data = [r"E:\AI\test_data\目标识别\storagetank\storagetank2.jpeg"]
    # params = {"object_name": "storagetank", "alg_model": "sat_hr-yolo_v11n-objdet-yolo_hbb", "conf_value": 0.45,
    #           "iou_value": 0.7,
    #           "image_data": image_data, "image_type": "path", "epsg_value": 4326}

    params["is_save_db"]=True
    task_id=snowflakeId()
    host = "192.168.3.80"
    dbengine = "django.db.backends.postgresql_psycopg2"
    port = 5432
    user = "postgres"
    password = "postgres123"
    db_name = "YNAI"
    table_name = "tt_ai_task_status"
    task_id_field="id"
    task_status_field="status"
    operate_db_info = {"db_conn_info": {"host": host, "port": port, "user": user, "password": password,
                                "db_name": db_name, "dbengine": dbengine},
               "update_sql": "update "+table_name+" set \"status\" = {status} , \"update_time\" = '{update_time}' where id =" + str(
                   task_id)}
    add_alg_task(task_id, table_name, task_id_field, task_status_field,"object_detection","sat_hr-yolo_v11n-objdet-yolo_hbb", operate_db_info)
    params["operate_db_info"] = operate_db_info
    params["bands"]=[[1,2,3]]
    #
    # db_conn_info={"db_conn_info": {"host": "192.168.3.80", "port": 5432, "user": "postgres", "password": "postgres123",
    #                   "db_name": "YNAI", "dbengine": "django.db.backends.postgresql_psycopg2"},
    #  "update_sql": "update tt_jimu_ai_task set status = 10 , update_time = {} where id =1797463448354816)}

    # # 测试路径普通图片（多张）
    # image_data = [r"E:\AI\test_data\目标识别\storagetank\storagetank2.jpeg",r"E:\AI\test_data\目标识别\storagetank\storagetank1.jpeg"]
    # params = {"object_name": "storagetank", "alg_model": "yolo_v11-objdet-yolo_hbb", "conf_value": 0.45, "iou_value": 0.7,
    #           "image_data": image_data, "image_type": "path"}
    # 测试base64普通图片
    # image_data = [
    #     ""]
    # params = {"object_name": "storagetank", "alg_model": "yolo_v11-objdet-yolo_hbb", "conf_value": 0.45, "iou_value": 0.7,
    #           "image_data": image_data, "image_type": "base_64"}
    #
    # 测试url普通图片
    # image_data = ["http://43.160.250.159:27615/storagetank1.jpeg","http://43.160.250.159:27615/storagetank2.jpeg"]
    # params = {"object_name": "storagetank", "alg_model": "yolo_v11-objdet-yolo_hbb", "conf_value": 0.45, "iou_value": 0.7,
    #           "image_data": image_data, "image_type": "url"}
    #
    res = image_detector.predictor_image(image_data, params)
    # 调用完成后，算法进度变成100
    update_alg_progress(100,operate_db_info)
    print(res)

    # # 测试视频检测器
    video_detector = VGISVideoPredictor(ai_service_url)
    # 测试url视频
    # video_data = ["http://43.160.250.159:27615/storagebank.mp4"]
    # params = {"object_name": "storagetank", "alg_model": "yolo_v8-objdet-yolo_hbb", "conf_value": 0.45, "iou_value": 0.7,
    #           "video_data": video_data, "video_type": "url"}
    #
    # # 测试stream视频
    # video_data = ["rtsp://192.168.31.79:8554/live/stream"]
    # params = {"object_name": "storagetank", "alg_model": "yolo_v11-objdet-yolo_hbb", "conf_value": 0.45, "iou_value": 0.7,
    #           "video_data": video_data, "video_type": "stream"}
    #
    # 测试路径视频E:\AI\test_data\目标识别\storagetank
    # video_data = ["E:\\AI\\test_data\\目标识别\\storagetank\\storagebank.mp4"]
    # params = {"object_name": "storagetank", "alg_model": "yolo_v11-objdet-yolo_hbb", "conf_value": 0.45,
    #           "iou_value": 0.7,
    #           "video_data": video_data, "video_type": "path"}
    #
    # # # 测试base64图像帧视频
    # frame1_image_data = ""
    # frame2_image_data = ""
    # video_data = [frame1_image_data, frame2_image_data]
    # params = {"object_name": "storagetank", "alg_model": "yolo_v11-objdet-yolo_hbb", "conf_value": 0.45, "iou_value": 0.7,
    #           "video_data": video_data, "video_type": "base_64"}
    # # # #
    #
    #
    # res = video_detector.predictor_video(video_data, params)
    # print(res)
